{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SMARTS\n",
    "This example demonstrates the use of SMARTS within a Jupyter notebook."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "wkR0YvENQni4"
   },
   "source": [
    "## Setup dependencies\n",
    "Install the base SMARTS dependencies needed for the example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "n1sQVa5mkMEA",
    "outputId": "dbfa09f5-0fd9-4a51-f728-fc62230f1b84"
   },
   "outputs": [],
   "source": [
    "# !git clone https://github.com/huawei-noah/SMARTS 2> /dev/null\n",
    "# !cd SMARTS && ls && pip install .[camera-obs]\n",
    "%pip install smarts[camera-obs]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Insert Colab content directory  \n",
    "This is only necessary if running this notebook on Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ubc6jEAqiCEq",
    "outputId": "16e9facc-86bc-4f72-a4f2-2d8bb2961435"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "sys.path.insert(0, Path(os.path.abspath(\"\")).parents[1])\n",
    "print(Path(os.path.abspath(\"\")) / \"SMARTS\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main\n",
    "The core example."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "s7UtcphinvNv"
   },
   "source": [
    "### Import Base Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HsnDWYx_ngdc"
   },
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "\n",
    "from smarts.zoo import registry\n",
    "from smarts.env.gymnasium.wrappers.episode_logger import EpisodeLogger\n",
    "from smarts.core.utils.episodes import episode_range\n",
    "from smarts.env.gymnasium.wrappers.single_agent import SingleAgent"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"smarts.env:hiway-v1\")\n",
    "env = SingleAgent(env)\n",
    "env: gym.Env = EpisodeLogger(env)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import zoo.policies.keep_lane_agent\n",
    "\n",
    "agent, _ = registry.make_agent(\"zoo.policies:keep-lane-agent-v0\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "LFoG7Z-FobPP"
   },
   "source": [
    "### Run the experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 501
    },
    "id": "BBWj3wbAso3J",
    "outputId": "63c1d665-b4a5-43b5-d926-b02b5bcce08f"
   },
   "outputs": [],
   "source": [
    "for episode in episode_range(max_steps=450):\n",
    "    observation = env.reset()\n",
    "    reward, terminated, truncated, info = None, False, False, None\n",
    "    while episode.continues(observation, reward, terminated, truncated, info):\n",
    "        action = agent.act(observation)\n",
    "        observation, reward, terminated, info = env.step(action)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Close the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "mock_demo.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "983a288b1deae516b5d1a3268f286490dc1a3bd215a042403a142fb3df5f8acd"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
